from typing import Dict, Any, overload, Literal, TypeVar, Union, Type
from abc import ABC, abstractmethod
import torch
import torch.nn as nn
from learners.registry import Registry

# Define a union type for all supported strategy types
StrategyType = Union[Literal["standard"], Literal["feature_dropout"]]


class ForwardStrategyInterface(ABC):
    """Protocol defining the interface for forward strategies."""

    @abstractmethod
    def forward_train(
        self, model: nn.Module, inputs: torch.Tensor, *args, **kwargs
    ) -> Dict[str, Any]:
        pass

    @abstractmethod
    def forward_test(
        self, model: nn.Module, inputs: torch.Tensor, *args, **kwargs
    ) -> Dict[str, Any]:
        pass

    @abstractmethod
    def forward_ca(
        self, model: nn.Module, inputs: torch.Tensor, *args, **kwargs
    ) -> Dict[str, Any]:
        pass


# Create a registry for forward strategies
FORWARD_STRATEGY_REGISTRY = Registry[ForwardStrategyInterface](
    "forward_strategy", ForwardStrategyInterface
)


@FORWARD_STRATEGY_REGISTRY.register("standard")
class StandardForwardStrategy(ForwardStrategyInterface):
    """Standard implementation of forward strategies.

    This implementation works with nn.Module and all its subclasses,
    passing appropriate flags to control the model's behavior.
    """

    def forward_train(
        self, model: nn.Module, inputs: torch.Tensor, *args, **kwargs
    ) -> Dict[str, Any]:
        """Execute model in training mode.

        Args:
            model: Model to execute
            inputs: Input tensors
            *args, **kwargs: Additional arguments passed to model

        Returns:
            Dictionary of model outputs
        """
        return model(inputs, training=True, *args, **kwargs)

    def forward_test(
        self, model: nn.Module, inputs: torch.Tensor, *args, **kwargs
    ) -> Dict[str, Any]:
        """Execute model in testing mode.

        Args:
            model: Model to execute
            inputs: Input tensors
            *args, **kwargs: Additional arguments passed to model

        Returns:
            Dictionary of model outputs
        """
        return model(inputs, training=False, *args, **kwargs)

    def forward_ca(
        self, model: nn.Module, inputs: torch.Tensor, *args, **kwargs
    ) -> Dict[str, Any]:
        """Execute model in classifier alignment mode.

        Args:
            model: Model to execute
            inputs: Input tensors
            *args, **kwargs: Additional arguments passed to model

        Returns:
            Dictionary of model outputs
        """
        return model(inputs, ca=True, *args, **kwargs)


# Example custom forward strategy
@FORWARD_STRATEGY_REGISTRY.register("feature_dropout")
class FeatureDropoutForwardStrategy(ForwardStrategyInterface):
    """Forward strategy that applies dropout to features during training.

    This demonstrates how to create a custom forward strategy that
    modifies the model behavior during different phases.
    """

    def __init__(self, dropout_rate: float = 0.5):
        """Initialize with configurable dropout rate.

        Args:
            dropout_rate: Probability of dropping features during training
        """
        self.dropout_rate = dropout_rate
        self.dropout = nn.Dropout(dropout_rate)

    def forward_train(
        self, model: nn.Module, inputs: torch.Tensor, *args, **kwargs
    ) -> Dict[str, Any]:
        """Apply dropout during training forward pass.

        Args:
            model: The model to execute
            inputs: Input tensors
            *args, **kwargs: Additional arguments

        Returns:
            Model outputs
        """
        # Apply dropout to inputs during training
        inputs = self.dropout(inputs)
        return model(inputs, training=True, *args, **kwargs)

    def forward_test(
        self, model: nn.Module, inputs: torch.Tensor, *args, **kwargs
    ) -> Dict[str, Any]:
        """Standard test mode without dropout.

        Args:
            model: The model to execute
            inputs: Input tensors
            *args, **kwargs: Additional arguments

        Returns:
            Model outputs
        """
        return model(inputs, training=False, *args, **kwargs)

    def forward_ca(
        self, model: nn.Module, inputs: torch.Tensor, *args, **kwargs
    ) -> Dict[str, Any]:
        """Classifier alignment forward pass with minimal dropout.

        Args:
            model: The model to execute
            inputs: Input tensors
            *args, **kwargs: Additional arguments

        Returns:
            Model outputs
        """
        # Use reduced dropout for classifier alignment
        if self.dropout_rate > 0.2:
            temp_dropout = nn.Dropout(0.2)  # Use lower dropout
            inputs = temp_dropout(inputs)
        return model(inputs, ca=True, *args, **kwargs)


# 确保这些类型被导出
__all__ = [
    "ForwardStrategyInterface",
    "StandardForwardStrategy",
    "FeatureDropoutForwardStrategy",
    "create_forward_strategy",
    "FORWARD_STRATEGY_REGISTRY",
    "StrategyType",
]


F = TypeVar("F", bound=ForwardStrategyInterface)


@overload
def create_forward_strategy(
    strategy_type: Literal["standard"],
) -> StandardForwardStrategy: ...


@overload
def create_forward_strategy(
    strategy_type: Literal["feature_dropout"], *, dropout_rate: float = 0.5
) -> FeatureDropoutForwardStrategy: ...


@overload
def create_forward_strategy(strategy_type: str, **kwargs) -> ForwardStrategyInterface: ...


def create_forward_strategy(strategy_type: str, **kwargs) -> ForwardStrategyInterface:
    """Factory function to create a forward strategy with precise type information."""
    if strategy_type == "standard":
        return FORWARD_STRATEGY_REGISTRY.create_typed(
            strategy_type, StandardForwardStrategy, **kwargs
        )
    elif strategy_type == "feature_dropout":
        return FORWARD_STRATEGY_REGISTRY.create_typed(
            strategy_type, FeatureDropoutForwardStrategy, **kwargs
        )
    else:
        # 通用情况
        return FORWARD_STRATEGY_REGISTRY.create(strategy_type, **kwargs)


# 泛型版本，当你需要更灵活的类型时
def create_typed_strategy(strategy_type: str, expected_type: Type[F], **kwargs) -> F:
    """创建指定类型的策略实例"""
    return FORWARD_STRATEGY_REGISTRY.create_typed(
        strategy_type, expected_type, **kwargs
    )
